use mark_dirty::DirtyMark;

use crate::scene::AsBytes;

use super::{buffer::UniformBuffer, Attribute};

pub struct Uniform<T: AsBytes> {
    data: DirtyMark<T>,
    buffer: UniformBuffer<T>,
}

impl<T: AsBytes + 'static> Attribute for Uniform<T> {
    fn layout(&self) -> &wgpu::BindGroupLayout {
        self.buffer.layout()
    }

    fn bind_group(&self) -> &wgpu::BindGroup {
        self.buffer.bind_group()
    }

    fn prepare(&mut self, ctx: &super::Context) {
        if self.data.is_dirty() {
            self.buffer.write(ctx, &self.data);
            self.data.unmark_dirty();
        }
    }
}

impl<T: AsBytes + Default> Uniform<T> {
    pub fn new_default(ctx: &super::Context, visibility: wgpu::ShaderStages) -> Self {
        let buffer = UniformBuffer::new(ctx, visibility);
        Self {
            data: DirtyMark::new(T::default()),
            buffer,
        }
    }
}

impl<T: AsBytes> Uniform<T> {
    pub fn new(ctx: &super::Context, visibility: wgpu::ShaderStages, data: T) -> Self {
        let buffer = UniformBuffer::new(ctx, visibility);
        Self {
            data: DirtyMark::new(data),
            buffer,
        }
    }

    pub fn set_data(&mut self, data: T) {
        *self.data = data;
    }

    pub fn data_mut(&mut self) -> &mut T {
        &mut self.data
    }
}
